package com.soFun.service;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * RAG服务，提供检索增强生成功能
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class RagService {

    private final ChatLanguageModel chatLanguageModel;
    private final EmbeddingModel embeddingModel;
    
    // 文档知识库映射表（每个知识库有唯一ID）
    private final Map<String, KnowledgeBase> knowledgeBases = new ConcurrentHashMap<>();
    
    /**
     * 创建知识库
     * @param name 知识库名称
     * @return 知识库ID
     */
    public String createKnowledgeBase(String name) {
        KnowledgeBase knowledgeBase = new KnowledgeBase(name);
        knowledgeBases.put(knowledgeBase.getId(), knowledgeBase);
        log.info("创建知识库：{}, ID: {}", name, knowledgeBase.getId());
        return knowledgeBase.getId();
    }
    
    /**
     * 添加文本到知识库
     * @param knowledgeBaseId 知识库ID
     * @param text 文本内容
     * @param metadata 元数据
     */
    public void addTextToKnowledgeBase(String knowledgeBaseId, String text, Map<String, String> metadata) {
        KnowledgeBase knowledgeBase = getKnowledgeBase(knowledgeBaseId);
        log.info("向知识库 {} 添加文本, 长度: {}", knowledgeBaseId, text.length());
        
        try {
            // 创建文档 - 尝试使用反射动态调用合适的方法
            Document document;
            try {
                // 使用反射获取适当的方法
                Class<?> documentClass = Document.class;
                try {
                    // 尝试使用静态方法from
                    Method fromMethod = documentClass.getMethod("from", String.class, Map.class);
                    document = (Document) fromMethod.invoke(null, text, metadata);
                } catch (NoSuchMethodException e) {
                    // 尝试使用构造函数
                    document = (Document) documentClass.getConstructor(String.class, Map.class).newInstance(text, metadata);
                }
            } catch (Exception ex) {
                log.error("无法创建Document对象: {}", ex.getMessage(), ex);
                throw new RuntimeException("无法创建Document对象", ex);
            }
            
            // 分割文档 - 使用反射确保方法兼容性
// Change the type of chunks to List<TextSegment>
            List<TextSegment> chunks;
            try {
                DocumentSplitter splitter = DocumentSplitters.recursive(500, 0);
                chunks = splitter.split(document);
                log.info("文本已分割为 {} 个片段", chunks.size());
            } catch (Exception e) {
                log.error("分割文档失败: {}", e.getMessage(), e);
                throw new RuntimeException("分割文档失败", e);
            }

// Adjust the subsequent code to work with TextSegment
            for (TextSegment chunk : chunks) {
                try {
                    // Create embedding vector
                    Embedding embedding = embeddingModel.embed(chunk.text()).content();

                    // Add to embedding store
                    knowledgeBase.getEmbeddingStore().add(embedding, chunk);
                } catch (Exception e) {
                    log.error("嵌入文档失败: {}", e.getMessage(), e);
                    throw new RuntimeException("文档嵌入处理失败", e);
                }
            }
            log.info("成功添加文本到知识库 {}", knowledgeBaseId);
        } catch (Exception e) {
            log.error("处理文档失败: {}", e.getMessage(), e);
            throw new RuntimeException("处理文档失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 从文件添加内容到知识库
     * @param knowledgeBaseId 知识库ID
     * @param filePath 文件路径
     * @throws IOException 如果文件读取失败
     */
    public void addFileToKnowledgeBase(String knowledgeBaseId, String filePath) throws IOException {
        Path path = Paths.get(filePath);
        log.info("从文件添加内容到知识库 {}, 文件路径: {}", knowledgeBaseId, path);
        String content = Files.readString(path);
        
        Map<String, String> metadata = new HashMap<>();
        metadata.put("source", path.getFileName().toString());
        
        addTextToKnowledgeBase(knowledgeBaseId, content, metadata);
    }
    
    /**
     * 使用知识库回答问题
     * @param knowledgeBaseId 知识库ID
     * @param question 问题
     * @return 回答
     */
    public String queryKnowledgeBase(String knowledgeBaseId, String question) {
        KnowledgeBase knowledgeBase = getKnowledgeBase(knowledgeBaseId);
        log.info("查询知识库 {}, 问题: {}", knowledgeBaseId, question);
        
        try {
            // 创建内容检索器 - 使用反射确保兼容性
            ContentRetriever contentRetriever;
            try {
                contentRetriever = EmbeddingStoreContentRetriever.builder()
                        .embeddingModel(embeddingModel)
                        .embeddingStore(knowledgeBase.getEmbeddingStore())
                        .maxResults(3)
                        .build();
            } catch (Exception e) {
                log.error("创建内容检索器失败: {}", e.getMessage(), e);
                throw new RuntimeException("创建内容检索器失败", e);
            }
            
            // 创建问答接口
            try {
                DocumentQueryService queryService = AiServices.builder(DocumentQueryService.class)
                        .chatLanguageModel(chatLanguageModel)
                        .contentRetriever(contentRetriever)
                        .build();
                
                // 查询
                String answer = queryService.answerQuestion(question);
                log.info("知识库查询完成，回答长度: {}", answer.length());
                return answer;
            } catch (Exception e) {
                log.error("创建查询服务失败: {}", e.getMessage(), e);
                throw new RuntimeException("创建查询服务失败", e);
            }
        } catch (Exception e) {
            log.error("知识库查询失败: {}", e.getMessage(), e);
            throw new RuntimeException("知识库查询处理失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 获取知识库信息
     * @param knowledgeBaseId 知识库ID
     * @return 知识库信息
     */
    public Map<String, Object> getKnowledgeBaseInfo(String knowledgeBaseId) {
        KnowledgeBase knowledgeBase = getKnowledgeBase(knowledgeBaseId);
        
        Map<String, Object> info = new HashMap<>();
        info.put("id", knowledgeBase.getId());
        info.put("name", knowledgeBase.getName());
        info.put("createdAt", knowledgeBase.getCreatedAt());
        info.put("documentCount", knowledgeBase.getDocumentCount());
        
        return info;
    }
    
    /**
     * 获取知识库对象
     * @param knowledgeBaseId 知识库ID
     * @return 知识库对象
     * @throws IllegalArgumentException 如果知识库不存在
     */
    private KnowledgeBase getKnowledgeBase(String knowledgeBaseId) {
        KnowledgeBase knowledgeBase = knowledgeBases.get(knowledgeBaseId);
        if (knowledgeBase == null) {
            log.warn("知识库不存在: {}", knowledgeBaseId);
            throw new IllegalArgumentException("知识库不存在，ID: " + knowledgeBaseId);
        }
        return knowledgeBase;
    }
    
    /**
     * 文档查询服务接口
     */
    interface DocumentQueryService {
        String answerQuestion(String question);
    }
    
    /**
     * 知识库类
     */
    @Getter
    static class KnowledgeBase {
        private final String id;
        private final String name;
        private final long createdAt;
        private final EmbeddingStore<TextSegment> embeddingStore;
        
        public KnowledgeBase(String name) {
            this.id = java.util.UUID.randomUUID().toString();
            this.name = name;
            this.createdAt = System.currentTimeMillis();
            this.embeddingStore = new InMemoryEmbeddingStore<>();
        }

        public int getDocumentCount() {
            // 使用try-catch处理可能的API不兼容问题
            try {
                return embeddingStore.addAll(null).size();
            } catch (NoSuchMethodError e) {
                // 如果size()方法不存在，尝试使用count()或其他可能的方法
                try {
                    Method countMethod = embeddingStore.getClass().getMethod("count");
                    return (int) countMethod.invoke(embeddingStore);
                } catch (Exception ex) {
                    log.error("获取文档数量失败: {}", ex.getMessage());
                    return 0; // 安全返回
                }
            }
        }
    }
} 